#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
脑电数据记录器

本文件用于从NeuroSky MindWave设备获取脑电信号数据并记录到日志文件中。
使用NeuroPy3库连接设备，通过回调函数机制实时获取各种脑电信号：
- 注意力(attention)和冥想(meditation)指数
- 原始脑电信号(rawValue)
- 各频段功率：delta, theta, lowAlpha, highAlpha, lowBeta, highBeta, lowGamma, midGamma
- 信号质量(poorSignal)和眨眼强度(blinkStrength)

原理：
1. 通过串口连接NeuroSky设备
2. 为每个数据类型设置回调函数
3. 当设备产生新数据时，回调函数被触发
4. 将数据以时间戳格式记录到CSV日志文件中
"""

import time
import csv
import os
from datetime import datetime
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from NeuroPy3 import NeuroPy

class EEGDataLogger:
    """脑电数据记录器类"""
    
    def __init__(self, port="COM5", baudrate=57600, log_file="eeg_data_wake.csv"):
        """
        初始化脑电数据记录器
        
        Args:
            port (str): 串口号，默认COM5
            baudrate (int): 波特率，默认57600
            log_file (str): 日志文件名，默认eeg_data.csv
        """
        self.port = port
        self.baudrate = baudrate
        self.log_file = log_file
        self.neuropy = None
        self.is_logging = False
        
        # 数据存储字典
        self.current_data = {
            'timestamp': None,
            'attention': 0,
            'meditation': 0,
            'rawValue': 0,
            'delta': 0,
            'theta': 0,
            'lowAlpha': 0,
            'highAlpha': 0,
            'lowBeta': 0,
            'highBeta': 0,
            'lowGamma': 0,
            'midGamma': 0,
            'poorSignal': 0,
            'blinkStrength': 0
        }
        
        # 初始化CSV文件
        self._init_csv_file()
    
    def _init_csv_file(self):
        """初始化CSV日志文件，写入表头"""
        if not os.path.exists(self.log_file):
            with open(self.log_file, 'w', newline='', encoding='utf-8') as f:
                writer = csv.writer(f)
                # 写入CSV表头
                headers = ['timestamp', 'attention', 'meditation', 'rawValue', 
                          'delta', 'theta', 'lowAlpha', 'highAlpha', 
                          'lowBeta', 'highBeta', 'lowGamma', 'midGamma', 
                          'poorSignal', 'blinkStrength']
                writer.writerow(headers)
    
    def _log_data_to_csv(self):
        """将当前数据记录到CSV文件"""
        if self.is_logging:
            self.current_data['timestamp'] = datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
            
            with open(self.log_file, 'a', newline='', encoding='utf-8') as f:
                writer = csv.writer(f)
                # 按照表头顺序写入数据
                row = [self.current_data[key] for key in 
                      ['timestamp', 'attention', 'meditation', 'rawValue', 
                       'delta', 'theta', 'lowAlpha', 'highAlpha', 
                       'lowBeta', 'highBeta', 'lowGamma', 'midGamma', 
                       'poorSignal', 'blinkStrength']]
                writer.writerow(row)
    
    # 回调函数定义
    def attention_callback(self, value):
        """注意力数据回调函数"""
        self.current_data['attention'] = value
        print(f"注意力: {value}")
        self._log_data_to_csv()
    
    def meditation_callback(self, value):
        """冥想数据回调函数"""
        self.current_data['meditation'] = value
        print(f"冥想: {value}")
        self._log_data_to_csv()
    
    def raw_value_callback(self, value):
        """原始脑电信号回调函数"""
        self.current_data['rawValue'] = value
        print(f"原始信号: {value}")
    
    def delta_callback(self, value):
        """Delta波回调函数"""
        self.current_data['delta'] = value
        print(f"Delta波: {value}")
        self._log_data_to_csv()
    
    def theta_callback(self, value):
        """Theta波回调函数"""
        self.current_data['theta'] = value
        print(f"Theta波: {value}")
    
    def low_alpha_callback(self, value):
        """低Alpha波回调函数"""
        self.current_data['lowAlpha'] = value
        print(f"低Alpha波: {value}")
    
    def high_alpha_callback(self, value):
        """高Alpha波回调函数"""
        self.current_data['highAlpha'] = value
        print(f"高Alpha波: {value}")
    
    def low_beta_callback(self, value):
        """低Beta波回调函数"""
        self.current_data['lowBeta'] = value
        print(f"低Beta波: {value}")
    
    def high_beta_callback(self, value):
        """高Beta波回调函数"""
        self.current_data['highBeta'] = value
        print(f"高Beta波: {value}")
    
    def low_gamma_callback(self, value):
        """低Gamma波回调函数"""
        self.current_data['lowGamma'] = value
        print(f"低Gamma波: {value}")
    
    def mid_gamma_callback(self, value):
        """中Gamma波回调函数"""
        self.current_data['midGamma'] = value
        print(f"中Gamma波: {value}")
    
    def poor_signal_callback(self, value):
        """信号质量回调函数"""
        self.current_data['poorSignal'] = value
        print(f"信号质量: {value}")
        self._log_data_to_csv()
    
    def blink_strength_callback(self, value):
        """眨眼强度回调函数"""
        self.current_data['blinkStrength'] = value
        print(f"眨眼强度: {value}")
    
    def start_logging(self):
        """开始记录脑电数据"""
        try:
            # 初始化NeuroPy对象
            self.neuropy = NeuroPy(self.port, self.baudrate)
            
            # 设置所有回调函数
            self.neuropy.setCallBack("attention", self.attention_callback)
            self.neuropy.setCallBack("meditation", self.meditation_callback)
            self.neuropy.setCallBack("rawValue", self.raw_value_callback)
            self.neuropy.setCallBack("delta", self.delta_callback)
            self.neuropy.setCallBack("theta", self.theta_callback)
            self.neuropy.setCallBack("lowAlpha", self.low_alpha_callback)
            self.neuropy.setCallBack("highAlpha", self.high_alpha_callback)
            self.neuropy.setCallBack("lowBeta", self.low_beta_callback)
            self.neuropy.setCallBack("highBeta", self.high_beta_callback)
            self.neuropy.setCallBack("lowGamma", self.low_gamma_callback)
            self.neuropy.setCallBack("midGamma", self.mid_gamma_callback)
            self.neuropy.setCallBack("poorSignal", self.poor_signal_callback)
            self.neuropy.setCallBack("blinkStrength", self.blink_strength_callback)
            
            # 开始数据采集
            self.is_logging = True
            self.neuropy.start()
            
            print(f"开始记录脑电数据到文件: {self.log_file}")
            print(f"串口: {self.port}, 波特率: {self.baudrate}")
            print("按 Ctrl+C 停止记录")
            
            # 保持程序运行
            while self.is_logging:
                time.sleep(1)
                
        except KeyboardInterrupt:
            print("\n用户中断，停止记录...")
            self.stop_logging()
        except Exception as e:
            print(f"错误: {e}")
            self.stop_logging()
    
    def stop_logging(self):
        """停止记录脑电数据"""
        self.is_logging = False
        if self.neuropy:
            self.neuropy.stop()
        print("脑电数据记录已停止")

def main():
    """主函数"""
    # 创建数据记录器实例
    logger = EEGDataLogger(port="COM5", log_file="eeg_data_wake.csv")
    
    # 开始记录数据
    logger.start_logging()

if __name__ == "__main__":
    main()